
'''
【复杂网络】社团划分结果评估指标：Q、ARI、NMI
https://bluebird.blog.csdn.net/article/details/110308488

一、模块度Q（Modularity）
Python：可以直接使用Community.modularity()包计算模块度。

二、兰德指数ARI（Adjusted Rand Index）
Python：使用sklearn.metrics.adjusted_rand_score(labels_true, labels_pred)计算。

三、标准化互信息NMI（Normalized Mutual Information）
假设对于N个样本点的两种标签划分为U 和 V. 熵为划分集的不准确性，定义如下:
Python: 使用sklearn.metrics.normalized_mutual_info_score(labels_true, labels_pred)计算。
'''

from sklearn import metrics
# import Community 
import numpy as np
import scipy.io as scio

def demo():
	a = np.random.randint(1,3,(1000))
	b = np.random.randint(1,3,(1000))
	print(a[:9])
	print(b[:9])

	# q = Community.modularity(a,b)
	ari = metrics.adjusted_rand_score(a, b)
	nmi = metrics.normalized_mutual_info_score(a, b)
	print(ari, nmi)


def txt2arr(txt):
	arr = []
	with open(txt, 'r') as f:#with语句自动调用close()方法
		for line in f.readlines():
			# print(line)
			num = int(line)
			arr.append(num)
	arr = np.array(arr)
	# print(arr.shape)
	# print(arr[:9])
	return arr

def my_adjusted_rand_score(labels_true, labels_pred):#metrics.adjusted_rand_score
    (tn, fp), (fn, tp) = metrics.pair_confusion_matrix(labels_true, labels_pred)
    # print((tn, fp), (fn, tp))
    tn = tn/1e6
    fp = fp/1e6
    fn = fn/1e6
    tp = tp/1e6
    # Special cases: empty data or full agreement
    if fn == 0 and fp == 0:
        return 1.0
    return 2. * (tp * tn - fn * fp) / ((tp + fn) * (fn + tn) + (tp + fp) * (fp + tn))

import argparse
parser = argparse.ArgumentParser(description="Argument")
parser.add_argument('--txt1', type=str, default='tb1m_anno.txt', help='')
parser.add_argument('--txt2', type=str, default='tb1m_out.txt', help='')
args = parser.parse_args()

if __name__ == '__main__':
	print('\n\n\n\t', 'Testing with ARI & NMI')
	print('*'*64)
		
	txt1 = args.txt1
	txt2 = args.txt2
	print('txt1:', txt1)
	print('txt2:', txt2)

	#RuntimeWarning: overflow encountered in longlong_scalars
	arr1 = txt2arr(txt1).reshape(-1)
	arr2 = txt2arr(txt2).reshape(-1)

	# arr1 = np.random.randint(1,3,(100000,))
	# arr2 = np.random.randint(1,3,(100000,))

	arr1 = arr1.astype(np.int64)
	arr2 = arr2.astype(np.int64)
	print('arr1:', arr1.shape, arr1.dtype)
	print('arr2:', arr2.shape, arr2.dtype)

	if(arr1.shape != arr2.shape):
		print('Array shape do not match!!!', arr1.shape, ' vs ', arr2.shape)
		dim = min(arr1.shape[0], arr2.shape[0])
		print('Use shape of:', dim)
		arr1 = arr1[:dim]
		arr2 = arr2[:dim]

	print('*'*64)
	ari = my_adjusted_rand_score(arr1, arr2)
	print("ari={:.8f}".format(ari))
	nmi = metrics.normalized_mutual_info_score(arr1, arr2)
	print("nmi={:.8f}".format(nmi))
	# acc = metrics.f1_score(arr1, arr2)
	# print(acc)
	print('*'*64)
	
	scio.savemat(txt2.replace('_out.txt', '_matlab.mat'), {'pred':arr2, 'true':arr1})